import numpy as np

class ManualKMeans:
    def __init__(self, n_clusters, max_iter=300, tol=1e-4, random_state=None):
        """
        手动实现的 KMeans 聚类算法
        :param n_clusters: 聚类数
        :param max_iter: 最大迭代次数
        :param tol: 收敛阈值，聚类中心的变化小于此值时停止
        :param random_state: 随机种子，保证结果可复现
        """
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.tol = tol
        self.random_state = random_state
        self.cluster_centers_ = None  # 聚类中心
        self.inertia_ = None  # 聚类惯性（WCSS）
        self.labels_ = None  # 每个点的聚类标签
        self.n_iter_ = 0  # 初始化迭代次数

    def fit(self, X):
        """
        对数据进行聚类
        :param X: 输入数据，形状为 (n_samples, n_features)
        """
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape

        # 随机初始化聚类中心
        initial_indices = np.random.choice(n_samples, self.n_clusters, replace=False)
        self.cluster_centers_ = X[initial_indices]

        for iteration in range(self.max_iter):
            self.n_iter_ = iteration + 1  # 更新迭代次数
            # Step 1: 计算每个点的距离并分配到最近的中心
            distances = self._compute_distances(X, self.cluster_centers_)
            labels = np.argmin(distances, axis=1)

            # Step 2: 重新计算每个聚类的中心
            new_centers = np.array([X[labels == i].mean(axis=0) for i in range(self.n_clusters)])

            # 检查是否收敛
            center_shift = np.linalg.norm(self.cluster_centers_ - new_centers, axis=1).max()
            self.cluster_centers_ = new_centers

            if center_shift < self.tol:
                break

        # 保存聚类结果
        self.labels_ = labels
        self.inertia_ = self._compute_inertia(X, labels)

    def predict(self, X):
        """
        为新数据点分配聚类标签
        :param X: 输入数据
        :return: 聚类标签
        """
        distances = self._compute_distances(X, self.cluster_centers_)
        return np.argmin(distances, axis=1)

    def fit_predict(self, X):
        """
        先拟合数据再预测标签
        :param X: 输入数据
        :return: 聚类标签
        """
        self.fit(X)
        return self.labels_

    def _compute_distances(self, X, centers):
        """
        计算样本点与聚类中心的距离
        :param X: 输入数据
        :param centers: 聚类中心
        :return: 每个样本点到各聚类中心的距离矩阵
        """
        return np.linalg.norm(X[:, np.newaxis, :] - centers, axis=2)

    def _compute_inertia(self, X, labels):
        """
        计算簇内平方和（WCSS）
        :param X: 输入数据
        :param labels: 聚类标签
        :return: WCSS 值
        """
        return sum(np.sum((X[labels == i] - center) ** 2) for i, center in enumerate(self.cluster_centers_))